#!/usr/bin/env python
# coding: utf-8
import pandas as pd
import yaml
from fbprophet import Prophet
import matplotlib.pyplot as plt


# In[ ]:

cfg = yaml.full_load(open("./config.yml", 'r',encoding='utf-8'))

hparams = cfg['HPARAMS']['PROPHET']

trainSet = pd.read_csv('./raw_data/train.csv')

traindata_len = len(trainSet)
print('train data: ',traindata_len)

# In[26]:

changepoint_prior_scale = 0.001
seasonality_prior_scale = 0.01
holidays_prior_scale = 0.1 #0.01
seasonality_mode = 'multiplicative' #'additive'
changepoint_range = 0.8
country = 'CN'

# Build DataFrame of local holidays

if hparams.get('HOLIDAYS', None) is None:
    local_holidays = None
else:
    holiday_dfs = []
    for holiday in hparams.get('HOLIDAYS', []):
        upper_w = 1
        if holiday[:5] == 'Covid':
            upper_w = 78  # 85
        elif holiday[:5] == 'Large':
                upper_w = 15
        else:
            upper_w = 1
        holiday_dfs.append(pd.DataFrame({
            'holiday': holiday,
            'ds': pd.to_datetime(hparams['HOLIDAYS'][holiday]),
            'lower_window': 0,
            'upper_window': upper_w}))
    local_holidays = pd.concat(holiday_dfs)

OUT_PATH ='./user_data/'
RESULT_PATH ='./prediction_result/'
outname = 'result_2.csv'

for step in range (0,3): #A\B厂 2个数据列分别建模预测，假设A、B无关联
    if step==0 :
        colname = 'A厂'
        df = trainSet.drop(['B厂'], axis=1)
        df.rename(columns={'日期': 'ds', colname: 'y'}, inplace=True)
    else:
        colname = 'B厂'
        df = trainSet.drop(['A厂'], axis=1)
        df.rename(columns={'日期': 'ds', colname: 'y'}, inplace=True)
        # 新冠时期处理
        if step == 1 :
            df.loc[(df.index >=749)  &  (df.index <= 839), colname] = ''

    df_ = df[:-60]  # 保留60个数据进行验证

    m = Prophet(holidays=local_holidays,# changepoints=['2020-01-18'], #yearly_seasonality=True,
                changepoint_prior_scale=changepoint_prior_scale,
                seasonality_prior_scale=seasonality_prior_scale, holidays_prior_scale=holidays_prior_scale,
                seasonality_mode=seasonality_mode, changepoint_range=changepoint_range)
    m.add_country_holidays(country_name=country)  # Add country-wide holidays

    #m = Prophet()

    m.fit(df_)

# In[28]:

    future = m.make_future_dataframe(periods=211)

    forecast = m.predict(future)
    #forecast = forecast.join(df, how='left',lsuffix='_id', rsuffix='_p')
    preds = forecast[1035:]

    if step == 0:
        pred = pd.DataFrame({'日期': preds['ds'], colname: preds['yhat']})
        #pred = pd.DataFrame({'日期': preds['ds'].tolist(), colname: preds['yhat'].tolist()})
        pred.loc[(pred.index -traindata_len <=33 ), colname] = pred[colname]*1.15
        pred.loc[(pred.index -traindata_len>=34)  &  (pred.index -traindata_len<= 52), colname] = pred[colname]*1.1
        pred.loc[(pred.index-traindata_len >=74)  &  (pred.index-traindata_len <= 84), colname] = pred[colname]*1.1
        pred.loc[(pred.index-traindata_len >=93)  &  (pred.index-traindata_len <= 103), colname] = pred[colname]*1.15
        pred.loc[(pred.index-traindata_len >=104) &  (pred.index -traindata_len<= 112), colname]=pred[colname]*0.92
        pred.loc[(pred.index-traindata_len >=113), colname]=pred[colname]*1.15

        pred[colname] = pred[colname].astype("int")
        pred_all = pred
    elif step == 1:
        trainSet.loc[(trainSet.index >=749)  &  (trainSet.index <= 839), colname] = forecast[749:840]
    else :
        pred_all[colname] = preds['yhat']
        pred_all.loc[(pred_all.index-traindata_len <= 6), colname] = pred_all[colname] * 0.95
        pred_all.loc[(pred_all.index-traindata_len >= 15) & (pred_all.index-traindata_len <= 33), colname] = pred_all[colname] * 1.05
        pred_all.loc[(pred_all.index-traindata_len >= 63) & (pred_all.index-traindata_len <= 92), colname] = pred_all[colname] * 0.92
        pred_all.loc[(pred_all.index-traindata_len >= 117) & (pred_all.index-traindata_len <= 122), colname] = pred_all[colname] * 1.3

        pred_all.loc[(pred_all.index-traindata_len >=123), colname]=pred_all[colname]*1.2
        pred_all[colname] = pred_all[colname].astype("int")

    print('step {}:{}'.format(step,outname) )
from datetime import datetime
pred_all['日期'] = pred_all['日期'].apply(lambda x: datetime.strftime(x,"%Y/%m/%d"))
pred_all.to_csv(RESULT_PATH+outname, index=False)
#融合模型结果
resSet1 = pd.read_csv(RESULT_PATH+'result_1.csv')
#pred_all = pred_all.join(resSet1, how='left', lsuffix='_1', rsuffix='_2')
pred_all2 = pd.merge(pred_all,resSet1, on='日期',suffixes=('_x', '_y')) #merge
#pred_all2 = pd.concat([pred_all,resSet1],axis=1, join='inner') #merge
pred_all2['A厂'] = pred_all2['A厂_x'] *0.5 + pred_all2['A厂_y']*0.5
pred_all2['B厂'] = pred_all2['B厂_x'] *0.5 + pred_all2['B厂_y']*0.5
pred_all2 = pd.DataFrame({'日期':pred_all2['日期'],'A厂':pred_all2['A厂'].astype("int"),'B厂':pred_all2['B厂'].astype("int")})
outname = 'result.csv'
pred_all2.to_csv(RESULT_PATH+outname, index=False)
print('模型融合完成!')
exit(1)
# In[37]:

'''
y0 = df['y']
print(y0.shape)


# In[44]:


x1 = forecast['ds']
y0 = df['y']
y1 = forecast['yhat']
y2 = forecast['yhat_lower']
y3 = forecast['yhat_upper']
pre = forecast['yhat']
plt.figure(figsize = (18,8))
plt.plot(x1[:-100],y0[:-100],label = 'original_data')
#plt.figure(figsize = (18,8))
#plt.plot(x1[:60],y0[:60],label = 'original_data')
#plt.plot(x1[-60:],y0[-60:],label = 'compared_data')
#plt.plot(x1[-60:],pre[-60:],label = 'predict_data')
plt.legend()
plt.show()


# In[13]:


fig2 = m.plot_components(forecast)
'''
